import tensorflow as tf
import keras

from tensorflow.keras.layers import Conv2D, Input, MaxPooling2D, Dropout, concatenate,UpSampling2D,AveragePooling2D,BatchNormalization,Activation,Lambda,SeparableConv2D


def ASPP(tensor):

    y_avgpool =  AveragePooling2D(pool_size=(2, 2),strides = 1,padding = 'same')(tensor)
    y_avgpool =  Conv2D(filters=128, kernel_size=1, padding='same',
                    kernel_initializer='he_normal', name='avgpool_1x1conv2d', use_bias=False)(y_avgpool)

    y_pool = Conv2D(filters=128, kernel_size=1, padding='same',
                    kernel_initializer='he_normal', name='pool_1x1conv2d', use_bias=False)(tensor)
    y_pool = BatchNormalization(name=f'bn_1')(y_pool)
    y_pool = Activation('relu', name=f'relu_1')(y_pool)

    y_1 = Conv2D(filters=128, kernel_size=3, dilation_rate=3, padding='same',
                 kernel_initializer='he_normal', name='ASPP_conv2d_d1', use_bias=False)(tensor)
    y_1 = BatchNormalization(name=f'bn_2')(y_1)
    y_1 = Activation('relu', name=f'relu_2')(y_1)


    y_6 = Conv2D(filters=128, kernel_size=3, dilation_rate=6, padding='same',
                 kernel_initializer='he_normal', name='ASPP_conv2d_d6', use_bias=False)(tensor)
    y_6 = BatchNormalization(name=f'bn_3')(y_6)
    y_6 = Activation('relu', name=f'relu_3')(y_6)

    y_12 = Conv2D(filters=128, kernel_size=3, dilation_rate=12, padding='same',
                  kernel_initializer='he_normal', name='ASPP_conv2d_d12', use_bias=False)(tensor)
    y_12 = BatchNormalization(name=f'bn_4')(y_12)
    y_12 = Activation('relu', name=f'relu_4')(y_12)

    y_18 = Conv2D(filters=128, kernel_size=3, dilation_rate=18, padding='same',
                  kernel_initializer='he_normal', name='ASPP_conv2d_d18', use_bias=False)(tensor)
    y_18 = BatchNormalization(name=f'bn_5')(y_18)
    y_18 = Activation('relu', name=f'relu_5')(y_18)


    y = concatenate([y_avgpool,y_pool, y_1, y_6, y_12, y_18])
    y_res = Conv2D(filters=512, kernel_size=3, dilation_rate=18, padding='same',
                  kernel_initializer='he_normal', name='res', use_bias=False)(tensor)

    y = Conv2D(filters=512, kernel_size=1, dilation_rate=1, padding='same',
               kernel_initializer='he_normal', name='ASPP_conv2d_final', use_bias=False)(y)
    y = BatchNormalization(name=f'bn_final')(y)
    y = Activation('relu', name=f'relu_final')(y)
    y = tf.keras.layers.add([y_res,y])
    return y